/**
 * ColorSensor
 *
 * Authors: Bruna and Bruno
 * Date: Jun, 09 2012
 **/
 
#include <Arduino.h>
#include <ColorSensor.h>

// ADJD-S311's I2C address, don't change
#define ADJD_S311_ADDRESS 0x74

// ADJD-S311's register list
#define CTRL 0x00
#define CONFIG 0x01

#define CAP_RED 0x06
#define CAP_GREEN 0x07
#define CAP_BLUE 0x08
#define CAP_CLEAR 0x09

#define INT_RED_LO 0xA
#define INT_RED_HI 0xB
#define INT_GREEN_LO 0xC
#define INT_GREEN_HI 0xD
#define INT_BLUE_LO 0xE
#define INT_BLUE_HI 0xF
#define INT_CLEAR_LO 0x10
#define INT_CLEAR_HI 0x11

#define DATA_RED_LO 0x40
#define DATA_RED_HI 0x41
#define DATA_GREEN_LO 0x42
#define DATA_GREEN_HI 0x43
#define DATA_BLUE_LO 0x44
#define DATA_BLUE_HI 0x45
#define DATA_CLEAR_LO 0x46
#define DATA_CLEAR_HI 0x47

#define OFFSET_RED 0x48
#define OFFSET_GREEN 0x49
#define OFFSET_BLUE 0x4A
#define OFFSET_CLEAR 0x4B

// can be found on datasheet
#define GSSR 0x01
#define GOFS 0x02
#define TOFS 0x01

// colors
#define RED 0
#define GREEN 1
#define BLUE 2
#define CLEAR 3

// color sensor led intensity
#define LED_ON 10

// constants used to determine color
#define ERR 2
#define TERM1 -1.8378770664

#define GreenUr 0.307107858715
#define GreenUg 0.362671676946
#define GreenI11 1976.73281635
#define GreenI12 902.105258658
#define GreenI22 4385.07851738
#define GreenDet 1.2731823088e-07

#define RedUr 0.549430496013
#define RedUg 0.248056353719
#define RedI11 2422.3343509
#define RedI12 3364.01263909
#define RedI22 8299.63442637
#define RedDet 1.13792718259e-07

#define LilacUr 0.350238818865
#define LilacUg 0.290040437694
#define LilacI11 1341.61299284
#define LilacI12 513.417812462
#define LilacI22 4177.15516539
#define LilacDet 1.87247410516e-07

#define YellowUr 0.442362319709
#define YellowUg 0.370131336452
#define YellowI11 3229.44265972
#define YellowI12 3520.01260886
#define YellowI22 6880.74501521
#define YellowDet 1.01724404473e-07

#define MIN_VALUE_FOR_ASSERTION 0
 
ColorSensor::ColorSensor(int _ledPin, int _samples) {
    ledPin = _ledPin;
    samples = _samples;
}
 
void ColorSensor::ledOn() {
    analogWrite(ledPin, LED_ON);
}
 
void ColorSensor::begin() {
    Wire.begin();
    delayMicroseconds(1);  // Wait for ADJD reset sequence
    reset();
}

void ColorSensor::reset() {
  colorCap[RED] = 15;
  colorCap[GREEN] = 11;
  colorCap[BLUE] = 5;
  colorCap[CLEAR] = 5;  
  
  colorInt[RED] = 4095;
  colorInt[GREEN] = 4095;
  colorInt[BLUE] = 4095;
  colorInt[CLEAR] = 4095;  
  
  init();
}

void ColorSensor::init() { 
  writeRegister(colorCap[RED] & 0xF, CAP_RED);
  writeRegister(colorCap[GREEN] & 0xF, CAP_GREEN);
  writeRegister(colorCap[BLUE] & 0xF, CAP_BLUE);
  writeRegister(colorCap[CLEAR] & 0xF, CAP_CLEAR);

  writeRegister((unsigned char)colorInt[RED], INT_RED_LO);
  writeRegister((unsigned char)((colorInt[RED] & 0x1FFF) >> 8), INT_RED_HI);
  writeRegister((unsigned char)colorInt[BLUE], INT_BLUE_LO);
  writeRegister((unsigned char)((colorInt[BLUE] & 0x1FFF) >> 8), INT_BLUE_HI);
  writeRegister((unsigned char)colorInt[GREEN], INT_GREEN_LO);
  writeRegister((unsigned char)((colorInt[GREEN] & 0x1FFF) >> 8), INT_GREEN_HI);
  writeRegister((unsigned char)colorInt[CLEAR], INT_CLEAR_LO);
  writeRegister((unsigned char)((colorInt[CLEAR] & 0x1FFF) >> 8), INT_CLEAR_HI);
}

Color ColorSensor::read() {
    int i;
    float term1 = TERM1 * samples;
    
    float redTerm3 = 0;
    float greenTerm3 = 0;
    float lilacTerm3 = 0;
    float yellowTerm3 = 0;

    float redTerm3_1 = 0;
    float greenTerm3_1 = 0;
    float lilacTerm3_1 = 0;
    float yellowTerm3_1 = 0;
    float redTerm3_2 = 0;
    float greenTerm3_2 = 0;
    float lilacTerm3_2 = 0;
    float yellowTerm3_2 = 0;
    float redTerm3_3 = 0;
    float greenTerm3_3 = 0;
    float lilacTerm3_3 = 0;
    float yellowTerm3_3 = 0;
    
    float div = samples / 2.0; 
    float redTerm2 = -div * log(RedDet);
    float greenTerm2 = -div * log(GreenDet);
    float lilacTerm2 = -div * log(LilacDet);
    float yellowTerm2 = -div * log(YellowDet);

    float red, green, lilac, yellow;

    Color maxColor;
    float maxValue;

    for (i=0; i<samples; i++) {
        getNormalizedRGB();
        redTerm3_1 += (normalizedRGB[RED] - RedUr) * (normalizedRGB[RED] - RedUr);
        redTerm3_2 += (normalizedRGB[RED] - RedUr) * (normalizedRGB[GREEN] - RedUg);
        redTerm3_3 += (normalizedRGB[GREEN] - RedUg) * (normalizedRGB[GREEN] - RedUg);
        
        greenTerm3_1 += (normalizedRGB[RED] - GreenUr) * (normalizedRGB[RED] - GreenUr);
        greenTerm3_2 += (normalizedRGB[RED] - GreenUr) * (normalizedRGB[GREEN] - GreenUg);
        greenTerm3_3 += (normalizedRGB[GREEN] - GreenUg) * (normalizedRGB[GREEN] - GreenUg);
        
        lilacTerm3_1 += (normalizedRGB[RED] - LilacUr) * (normalizedRGB[RED] - LilacUr);
        lilacTerm3_2 += (normalizedRGB[RED] - LilacUr) * (normalizedRGB[GREEN] - LilacUg);
        lilacTerm3_3 += (normalizedRGB[GREEN] - LilacUg) * (normalizedRGB[GREEN] - LilacUg);
        
        yellowTerm3_1 += (normalizedRGB[RED] - YellowUr) * (normalizedRGB[RED] - YellowUr);
        yellowTerm3_2 += (normalizedRGB[RED] - YellowUr) * (normalizedRGB[GREEN] - YellowUg);
        yellowTerm3_3 += (normalizedRGB[GREEN] - YellowUg) * (normalizedRGB[GREEN] - YellowUg);
    }
  
    redTerm3 = redTerm3_1 * RedI11 + 2 * redTerm3_2 * RedI12 + redTerm3_3 * RedI22;
    greenTerm3 = greenTerm3_1 * GreenI11 + 2 * greenTerm3_2 * GreenI12 + greenTerm3_3 * GreenI22;
    lilacTerm3 = lilacTerm3_1 * LilacI11 + 2 * lilacTerm3_2 * LilacI12 + lilacTerm3_3 * LilacI22;
    yellowTerm3 = yellowTerm3_1 * YellowI11 + 2 * yellowTerm3_2 * YellowI12 + yellowTerm3_3 * YellowI22;
  
    redTerm3 /= 2;
    greenTerm3 /= 2;
    lilacTerm3 /= 2;
    yellowTerm3 /= 2;

    red = term1 + redTerm2 - redTerm3;
    green = term1 + greenTerm2 - greenTerm3;
    lilac = term1 + lilacTerm2 - lilacTerm3;
    yellow = term1 + yellowTerm2 - yellowTerm3;
    
    maxValue = MIN_VALUE_FOR_ASSERTION;
    maxColor = _NONE;

    if (maxValue < red) {
        maxValue = red;
        maxColor = _RED;
    }
    
    if (maxValue < green) {
        maxValue = green;
        maxColor = _GREEN;
    }

    if (maxValue < lilac) {
        maxValue = lilac;
        maxColor = _LILAC;
    }

    if (maxValue < yellow) {
        maxValue = yellow;
        maxColor = _YELLOW;
    }

    return maxColor;
}

void ColorSensor::getNormalizedRGB() {
  getRGBC();

  float sum = (colorData[RED] + colorData[GREEN]+ colorData[BLUE]);
  normalizedRGB[RED] = (colorData[RED] / sum);
  normalizedRGB[GREEN] = (colorData[GREEN] / sum);
}

/** 
 * This function reads all of the ADJD-S311's data registers and stores them into colorData[].
 **/
void ColorSensor::getRGBC() {
  measure();
  
  colorData[RED] = readRegisterInt(DATA_RED_LO);
  colorData[GREEN] = readRegisterInt(DATA_GREEN_LO);
  colorData[BLUE] = readRegisterInt(DATA_BLUE_LO);
  colorData[CLEAR] = readRegisterInt(DATA_CLEAR_LO);
}

void ColorSensor::measure() {  
  writeRegister(GSSR, CTRL); // start sensing
  while(readRegister(CTRL) != 0) {} // waiting for a result
}

void ColorSensor::calibrate() {
    reset();
    calibrateColor();
    calibrateClear();
    calibrateCapacitors();
}

/** 
 * This function calibrates the clear integration registers of the ADJD-S311.
 **/
void ColorSensor::calibrateClear() {
  int gainFound = 0;
  int upperBox = 4096;
  int lowerBox = 0;
  int half;
  
  while (!gainFound)
  {
    half = ( (upperBox-lowerBox) / 2 ) + lowerBox;
    //no further halfing possbile
    if (half == lowerBox)
      gainFound = 1;
    else 
    {
      writeRegisterInt(INT_CLEAR_LO, half);
      measure();
      int halfValue = readRegisterInt(DATA_CLEAR_LO);

      if (halfValue>1000)
        upperBox = half;
      else if (halfValue<1000)
        lowerBox = half;
      else
        gainFound = 1;
    }
  }
}

/** 
 * This function clalibrates the RG and B integration registers.
 **/
void ColorSensor::calibrateColor() {
  int red, green, blue;
  int gainFound = 0;
  int upperBox = 4096;
  int lowerBox = 0;
  int half;
  
  while (!gainFound)
  {
    half = ( (upperBox-lowerBox) / 2 ) + lowerBox;
    //no further halfing possbile
    if (half == lowerBox)
    {
      gainFound = 1;
    }
    else
    {
      writeRegisterInt(INT_RED_LO, half);
      writeRegisterInt(INT_GREEN_LO, half);
      writeRegisterInt(INT_BLUE_LO, half);
      
      red = green = blue = 0;
      for (int i=0; i<4 ;i++)
      {
        measure();
        red   += readRegisterInt(DATA_RED_LO);
        green += readRegisterInt(DATA_GREEN_LO);
        blue  += readRegisterInt(DATA_BLUE_LO);
      }
      red   /= 4;
      green /= 4;
      blue  /= 4;
    
      int halfValue = 0;

      halfValue = max(halfValue, red);
      halfValue = max(halfValue, green);
      halfValue = max(halfValue, blue);

      if (halfValue > 1000)
      {
        upperBox = half;
      }
      else if (halfValue < 1000)
      {
        lowerBox = half;
      }
      else
      {
        gainFound = 1;
      }
    }
  }
}

/** 
 * This function calibrates each of the RGB and C capacitor registers.
 **/
void ColorSensor::calibrateCapacitors() {
  int calibrationRed = 0;
  int calibrationBlue = 0;
  int calibrationGreen = 0;
  int calibrated = 0;

  //need to store detect better calibration
  int oldDiff = 5000;

  while (!calibrated)
  {
    // sensor gain setting (Avago app note 5330)
    // CAPs are 4bit (higher value will result in lower output)
    writeRegister(calibrationRed, CAP_RED);
    writeRegister(calibrationGreen, CAP_GREEN);
    writeRegister(calibrationBlue, CAP_BLUE);

    int colorGain = readRegisterInt(INT_RED_LO);
    writeRegisterInt(INT_RED_LO, colorGain);
    writeRegisterInt(INT_GREEN_LO, colorGain);
    writeRegisterInt(INT_BLUE_LO, colorGain);

    int maxRead = 0;
    int minRead = 4096;
    int red   = 0;
    int green = 0;
    int blue  = 0;
    
    for (int i=0; i<4 ;i++)
    {
      measure();
      red   += readRegisterInt(DATA_RED_LO);
      green += readRegisterInt(DATA_GREEN_LO);
      blue  += readRegisterInt(DATA_BLUE_LO);
    }
    red   /= 4;
    green /= 4;
    blue  /= 4;

    maxRead = max(maxRead, red);
    maxRead = max(maxRead, green);
    maxRead = max(maxRead, blue);

    minRead = min(minRead, red);
    minRead = min(minRead, green);
    minRead = min(minRead, blue);

    int diff = maxRead - minRead;

    if (abs(oldDiff - diff) >= ERR)
    {
      if ((maxRead == red) && (calibrationRed < 15))
        calibrationRed++;
      else if ((maxRead == green) && (calibrationGreen < 15))
        calibrationGreen++;
      else if ((maxRead == blue) && (calibrationBlue < 15))
        calibrationBlue++;
    }
    else
      calibrated = 1;
      
    oldDiff = diff;
  }
}

/**
 * I2C functions
 * Write a byte of data to a specific ADJD-S311 address
 **/
void ColorSensor::writeRegister(unsigned char data, unsigned char address) {
  Wire.beginTransmission(ADJD_S311_ADDRESS);
  Wire.write(address);
  Wire.write(data);
  Wire.endTransmission();
}

/**
 * I2C functions
 * Read a byte of data from ADJD-S311 address
 **/
unsigned char ColorSensor::readRegister(unsigned char address) {
  unsigned char data;
  
  Wire.beginTransmission(ADJD_S311_ADDRESS);
  Wire.write(address);
  Wire.endTransmission();
  
  Wire.requestFrom(ADJD_S311_ADDRESS, 1);
  while (!Wire.available()) {}  // wait till we can get data
  
  return Wire.read();
}

/**
 * I2C functions
 * This function writes a 12-bit value to the LO and HI integration registers
 **/
void ColorSensor::writeRegisterInt(int address, int gain) {
  if (gain < 4096) 
  {
    byte msb = gain >> 8;
    byte lsb = gain;

    writeRegister(lsb, address);
    writeRegister(msb, address + 1);
  }
}

/**
 * I2C functions
 * Reads two bytes of data from ADJD-S311 address and address + 1
 **/
int ColorSensor::readRegisterInt(unsigned char address) {
  return readRegister(address) + (readRegister(address+1) << 8);
}